import logging
from typing import List

from langchain_core.documents import Document
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.constants import START, END
from langgraph.graph import StateGraph
from langgraph.prebuilt.chat_agent_executor import AgentState, create_react_agent

import util
from app.agent.state import FlightServiceState
from app.services.booking_service import BookingService
from app.services.flight_service import FlightService
from app.services.passenger_service import PassengerService
from app.retriever.know_retriever import FaqRetriever
from app.config.database_config import PgManager
from app.config.llm_client import create_chat_client
from app.tools.tool_set import tools

logger=logging.getLogger(__name__)

class FaqAgentStates(FlightServiceState):
    """子图状态"""
    passenger_id: str
    pass


class FaqAgentState(AgentState):
    pass

class FaqAgent:

    def __init__(self, pg_manager: PgManager):
        self.llm = create_chat_client()
        self.pg_manager = pg_manager
        self.booking_service = BookingService(next(pg_manager.get_db_static()))
        self.flight_service = FlightService(next(pg_manager.get_db_static()))
        self.passenger_service = PassengerService(next(pg_manager.get_db_static()))
        self.retriever = FaqRetriever()
        self.workflow = self.create_workflow()

    def _chat(self, state: FlightServiceState):
        """聊天节点"""
        agent = create_react_agent(model=self.llm,
                                   tools=tools,
                                   debug=True,
                                   state_schema=FaqAgentState
                                   )

        system_prompt = """
        你是一个专业的航班服务助手，能够帮助用户解答各种航班相关的问题并提供服务。

        你的主要功能包括：
        1. 回答航班相关的常见问题，如退票政策、行李规定、值机流程等
        2. 查询航班信息（需要调用工具）
        3. 帮助用户预订机票（需要调用工具）

        请根据用户的提问、对话历史以及提供的话术来给予友好的回复，语气要礼貌且专业。
        请结合以下提供的话术上下文，来给与用户回复，如果用户提问的内容不够清晰，请从其他角度进行引导，直到用户明确表达意图。

        话术上下文：
        {context}

        你可以调用相关工具来完成任务。
        """

        # 注意这里的history实际上已经包含了用户的提问，所以不需要再添加用户提问
        history_msgs = util.trunc_msgs(state.messages)
        faq_retriever = util.create_context_retriever(self.llm, self.retriever.retriever)
        faq_docs: List[Document] = faq_retriever.invoke(
            {"input": state.input})
        logger.warning(faq_docs)
        msgs = [SystemMessage(
            content=system_prompt.format(context=util.combine(faq_docs)))]+state.messages
        agent_res = agent.invoke({
            "messages": msgs,
        })
        logger.info(f"agent回复: {agent_res['messages'][-1].content}")
        return {
            "messages": [agent_res["messages"][-1]],
        }

    def create_workflow(self):
        """创建售前咨询子图工作流"""
        workflow = StateGraph(FaqAgentStates)
        workflow.add_node("chat", self._chat)
        workflow.add_edge(START, "chat")
        workflow.add_edge("chat", END)
        return workflow.compile()

    def run(self, question: str, session_id: str) -> FaqAgentStates:
        # 初始化输入
        initial_state = FaqAgentStates(
            input=question,
            passenger_id=session_id,
            messages=[
                HumanMessage(content=question)
            ],
            current_intent=None,
            error_message=None,
            human_approval_needed=None,
            human_approval_status=None,
            human_handoff_requested=None,

        )

        # 这里可以设置一个错误捕获机制，当深度超过设置时，使用人工介入
        config: RunnableConfig = {
            "configurable": {
                "thread_id": session_id
            },
            "recursion_limit": 10,
        }
        return self.workflow.invoke(initial_state, config=config)

if __name__ == "__main__":
    res = FaqAgent(PgManager()).run("我的航班预定号是BK20250805559766，查询一下信息", "1")
    print(res)